# -*- coding: utf-8 -*-
# @Time    : 2021/7/7 14:46
# @Author  : LuoTianHang


# ####################model.py 说明##########################
# this script is used to construct the model
from torch import nn


class ResBlock(nn.Module):
    def __init__(self, in_channel, out_channel=32):
        super().__init__()
        self.conv = nn.Conv2d(in_channel, out_channel, [3, 3], padding=1, bias=False, padding_mode="zeros")
        self.bn1 = nn.BatchNorm2d(out_channel)
        self.conv1 = nn.Conv2d(out_channel, out_channel, [3, 3], padding=1, bias=False, padding_mode="zeros")
        self.bn2 = nn.BatchNorm2d(out_channel)
        # self.leaky_relu = nn.LeakyReLU(inplace=True)
        # self.leaky_relu=nn.Identity()
        # self.leaky_relu=nn.Tanh()
        # self.leaky_relu=nn.Hardswish()
        self.leaky_relu=nn.ReLU(inplace=True)

    def forward(self, inputs):
        x = self.bn2(self.conv1(self.leaky_relu(self.bn1(self.conv1(inputs)))))
        return x + inputs


class UnetGenerator(nn.Module):
    def __init__(self, channel=32, num_blocks=4):
        super().__init__()

        self.conv = nn.Conv2d(3, channel, [3, 3], padding=1, bias=False, padding_mode="zeros")  # same 256,256
        self.bn = nn.InstanceNorm2d(channel)
        self.conv1 = nn.Conv2d(channel, channel, [3, 3], stride=2, padding=1, bias=False,
                               padding_mode="zeros")  # same 128,128
        self.bn1 = nn.InstanceNorm2d(channel)
        self.conv2 = nn.Conv2d(channel, channel * 2, [3, 3], padding=1, bias=False, padding_mode="zeros")  # 128,128
        self.bn2 = nn.InstanceNorm2d(channel * 2)
        self.conv3 = nn.Conv2d(channel * 2, channel * 2, [3, 3], stride=2, padding=1, bias=False,
                               padding_mode="zeros")  # 64,64
        self.bn3 = nn.InstanceNorm2d(channel * 2)
        self.conv4 = nn.Conv2d(channel * 2, channel * 4, [3, 3], padding=1, bias=False, padding_mode="zeros")  # 64,64
        self.bn4 = nn.InstanceNorm2d(channel * 4)
        self.resblock = nn.Sequential(*[ResBlock(channel * 4, channel * 4) for i in range(num_blocks)])

        self.conv5 = nn.Conv2d(channel * 4, channel * 2, [3, 3], padding=1, bias=False, padding_mode="zeros")  # 64,64
        self.bn5 = nn.InstanceNorm2d(channel * 2)
        self.conv6 = nn.Conv2d(channel * 2, channel * 2, [3, 3], padding=1, bias=False, padding_mode="zeros")  # 64,64
        self.bn6 = nn.InstanceNorm2d(channel * 2)
        self.conv7 = nn.Conv2d(channel * 2, channel, [3, 3], padding=1, bias=False, padding_mode="zeros")  # 64,64
        self.bn7 = nn.InstanceNorm2d(channel)
        self.conv8 = nn.Conv2d(channel, channel, [3, 3], padding=1, bias=False, padding_mode="zeros")  # 64,64
        self.bn8 = nn.InstanceNorm2d(channel)
        self.conv9 = nn.Conv2d(channel, channel, [3, 3], padding=1, bias=False, padding_mode="zeros")  # 64,64
        self.bn9 = nn.BatchNorm2d(channel)

        self.leak_relu = nn.LeakyReLU(inplace=True)
        self.leak_relu=nn.ReLU(inplace=True)
        self.upsample = nn.UpsamplingBilinear2d(scale_factor=2)
        self.conv10 = nn.Conv2d(channel, 3, [3, 3], padding=1, bias=False, padding_mode="zeros")  # 64,64
        self.bn10=nn.InstanceNorm2d(3)
        # self.conv11 = nn.Conv2d(channel, 3, [3, 3], padding=1, bias=False, padding_mode="zeros")  # 64,64
        self.act = nn.Tanh()

    def forward(self, inputs):
        x0 = self.conv(inputs)
        x0 = self.bn(x0)
        x0 = self.leak_relu(x0)  # 256, 256, 32

        x1 = self.conv1(x0)
        x1 = self.bn1(x1)
        x1 = self.leak_relu(x1)
        x1 = self.conv2(x1)
        x1 = self.bn2(x1)
        x1 = self.leak_relu(x1)  # 128, 128, 64

        x2 = self.conv3(x1)
        x2 = self.bn3(x2)
        x2 = self.leak_relu(x2)
        x2 = self.conv4(x2)
        x2 = self.bn4(x2)
        x2 = self.leak_relu(x2)  # 64, 64, 128

        x2 = self.resblock(x2)  # 64, 64, 128
        x2 = self.conv5(x2)
        x2 = self.bn5(x2)
        x2 = self.leak_relu(x2)  # 64, 64, 64

        x3 = self.upsample(x2)
        # x3 = self.conv6(x3 + x1)
        x3=self.conv6(x3)
        x3 = self.bn6(x3)
        x3 = self.leak_relu(x3)
        x3 = self.conv7(x3)
        x3 = self.bn7(x3)
        x3 = self.leak_relu(x3)  # 128, 128, 32

        x4 = self.upsample(x3)
        # x4 = self.conv8(x4 + x0)
        x4=self.conv8(x4)
        x4 = self.bn8(x4)
        x4 = self.leak_relu(x4)
        x4 = self.conv9(x4)  # 256, 256, 32
        x4 = self.bn9(x4)
        x4 = self.conv10(x4)
        x4=self.bn10(x4)
        result = self.act(x4)

        return result
